#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
目录表结构更新脚本

为directories表添加团队分配相关字段，建立与teams表的外键关联。
"""

import os
import sqlite3
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)

class DirectorySchemaUpdater:
    """目录表结构更新器"""

    def __init__(self, db_path: str = None):
        """初始化更新器"""
        if db_path is None:
            # 默认数据库路径
            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            db_path = os.path.join(project_root, 'backend', 'data', 'uma_insight.db')

        self.db_path = db_path

    def connect_database(self) -> sqlite3.Connection:
        """连接数据库"""
        try:
            conn = sqlite3.connect(self.db_path)
            # 启用外键约束
            conn.execute("PRAGMA foreign_keys = ON")
            return conn
        except Exception as e:
            logger.error(f"连接数据库失败: {e}")
            raise

    def check_table_exists(self, conn: sqlite3.Connection, table_name: str) -> bool:
        """检查表是否存在"""
        cursor = conn.cursor()
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name=?
        """, (table_name,))
        return cursor.fetchone() is not None

    def check_column_exists(self, conn: sqlite3.Connection, table_name: str, column_name: str) -> bool:
        """检查列是否存在"""
        cursor = conn.cursor()
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [row[1] for row in cursor.fetchall()]
        return column_name in columns

    def backup_table(self, conn: sqlite3.Connection, table_name: str):
        """备份表数据"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_table_name = f"{table_name}_backup_{timestamp}"

        logger.info(f"创建表备份: {backup_table_name}")
        conn.execute(f"CREATE TABLE {backup_table_name} AS SELECT * FROM {table_name}")
        conn.commit()

    def update_directories_schema(self, conn: sqlite3.Connection):
        """更新directories表结构"""
        logger.info("开始更新directories表结构...")

        # 检查表是否存在
        if not self.check_table_exists(conn, 'directories'):
            logger.error("directories表不存在，请先创建该表")
            return False

        # 检查teams表是否存在
        if not self.check_table_exists(conn, 'teams'):
            logger.error("teams表不存在，无法建立外键关联")
            return False

        # 备份现有数据
        self.backup_table(conn, 'directories')

        # 添加team_id字段（外键）
        if not self.check_column_exists(conn, 'directories', 'team_id'):
            logger.info("添加team_id字段...")
            conn.execute("""
                ALTER TABLE directories
                ADD COLUMN team_id TEXT
            """)
        else:
            logger.info("team_id字段已存在")

        # 添加assignment_confidence字段（分配置信度）
        if not self.check_column_exists(conn, 'directories', 'assignment_confidence'):
            logger.info("添加assignment_confidence字段...")
            conn.execute("""
                ALTER TABLE directories
                ADD COLUMN assignment_confidence INTEGER DEFAULT 0
            """)
        else:
            logger.info("assignment_confidence字段已存在")

        # 添加assigned_at字段（分配时间）
        if not self.check_column_exists(conn, 'directories', 'assigned_at'):
            logger.info("添加assigned_at字段...")
            conn.execute("""
                ALTER TABLE directories
                ADD COLUMN assigned_at DATETIME DEFAULT NULL
            """)
        else:
            logger.info("assigned_at字段已存在")

        # 添加assignment_method字段（分配方法）
        if not self.check_column_exists(conn, 'directories', 'assignment_method'):
            logger.info("添加assignment_method字段...")
            conn.execute("""
                ALTER TABLE directories
                ADD COLUMN assignment_method TEXT DEFAULT NULL
            """)
        else:
            logger.info("assignment_method字段已存在")

        # 创建外键约束
        logger.info("创建外键约束...")
        try:
            # SQLite不支持直接添加外键约束，需要重建表
            cursor = conn.cursor()
            cursor.execute("""
                SELECT sql FROM sqlite_master
                WHERE type='table' AND name='directories'
            """)
            original_sql = cursor.fetchone()[0]

            # 创建新的表结构
            new_table_sql = original_sql.replace(
                "CREATE TABLE directories",
                "CREATE TABLE directories_new"
            )

            # 添加外键约束
            if "FOREIGN KEY" not in new_table_sql:
                new_table_sql = new_table_sql.rstrip(");") + """,
                FOREIGN KEY (team_id) REFERENCES teams(id) ON DELETE SET NULL
                );"""

            # 创建新表
            logger.info("创建新的directories表结构...")
            conn.execute(new_table_sql)

            # 复制数据
            logger.info("复制现有数据到新表...")
            conn.execute("""
                INSERT INTO directories_new
                SELECT * FROM directories
            """)

            # 删除原表
            conn.execute("DROP TABLE directories")

            # 重命名新表
            conn.execute("ALTER TABLE directories_new RENAME TO directories")

            conn.commit()
            logger.info("外键约束创建成功")

        except Exception as e:
            logger.warning(f"创建外键约束时出现问题: {e}")
            logger.info("表结构已更新，但外键约束可能需要手动检查")

        logger.info("directories表结构更新完成！")
        return True

    def create_indexes(self, conn: sqlite3.Connection):
        """创建索引优化查询性能"""
        logger.info("创建性能优化索引...")

        indexes = [
            ("idx_directories_team_id", "CREATE INDEX IF NOT EXISTS idx_directories_team_id ON directories(team_id)"),
            ("idx_directories_confidence", "CREATE INDEX IF NOT EXISTS idx_directories_confidence ON directories(assignment_confidence)"),
            ("idx_directories_assigned_at", "CREATE INDEX IF NOT EXISTS idx_directories_assigned_at ON directories(assigned_at)"),
        ]

        for index_name, sql in indexes:
            try:
                conn.execute(sql)
                logger.info(f"索引 {index_name} 创建成功")
            except Exception as e:
                logger.warning(f"索引 {index_name} 创建失败: {e}")

        conn.commit()
        logger.info("索引创建完成")

    def verify_schema(self, conn: sqlite3.Connection):
        """验证更新后的表结构"""
        logger.info("验证表结构...")

        cursor = conn.cursor()

        # 检查表是否存在
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='directories'
        """)
        if not cursor.fetchone():
            logger.error("directories表不存在")
            return False

        # 检查列是否存在
        cursor.execute("PRAGMA table_info(directories)")
        columns = [row[1] for row in cursor.fetchall()]

        required_columns = ['team_id', 'assignment_confidence', 'assigned_at', 'assignment_method']
        missing_columns = [col for col in required_columns if col not in columns]

        if missing_columns:
            logger.error(f"缺少必需的列: {missing_columns}")
            return False

        logger.info("表结构验证通过")
        return True

    def show_table_info(self, conn: sqlite3.Connection):
        """显示表信息"""
        logger.info("\n=== directories表结构信息 ===")

        cursor = conn.cursor()
        cursor.execute("PRAGMA table_info(directories)")

        logger.info("列信息:")
        for row in cursor.fetchall():
            cid, name, type_, notnull, default, pk = row
            null_flag = "NOT NULL" if notnull else "NULL"
            pk_flag = "PRIMARY KEY" if pk else ""
            logger.info(f"  {name}: {type_} {null_flag} {pk_flag} (默认值: {default})")

        # 显示统计信息
        cursor.execute("SELECT COUNT(*) FROM directories")
        total_dirs = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM directories WHERE team_id IS NOT NULL")
        assigned_dirs = cursor.fetchone()[0]

        logger.info(f"\n=== 统计信息 ===")
        logger.info(f"总目录数: {total_dirs}")
        logger.info(f"已分配团队数: {assigned_dirs}")
        logger.info(f"分配覆盖率: {(assigned_dirs/total_dirs*100):.2f}%" if total_dirs > 0 else "0.00%")

    def run_update(self):
        """执行完整的更新流程"""
        logger.info("开始目录表结构更新流程")

        try:
            conn = self.connect_database()

            # 更新表结构
            if not self.update_directories_schema(conn):
                logger.error("表结构更新失败")
                return False

            # 创建索引
            self.create_indexes(conn)

            # 验证结构
            if not self.verify_schema(conn):
                logger.error("表结构验证失败")
                return False

            # 显示信息
            self.show_table_info(conn)

            logger.info("数据库表结构更新完成！")
            return True

        except Exception as e:
            logger.error(f"更新过程出错: {e}")
            return False
        finally:
            conn.close()


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='目录表结构更新工具')
    parser.add_argument('--db-path', help='数据库文件路径')
    parser.add_argument('--backup-only', action='store_true', help='仅创建备份')

    args = parser.parse_args()

    try:
        updater = DirectorySchemaUpdater(args.db_path)

        if args.backup_only:
            logger.info("仅创建备份模式")
            conn = updater.connect_database()
            updater.backup_table(conn, 'directories')
            conn.close()
            logger.info("备份创建完成")
        else:
            success = updater.run_update()
            if success:
                logger.info("表结构更新任务完成")
                return 0
            else:
                logger.error("表结构更新任务失败")
                return 1

    except Exception as e:
        logger.error(f"执行过程出错: {e}")
        return 1

    return 0


if __name__ == "__main__":
    exit(main())